from typing import List, Dict, Any, Optional

from langchain.chains.base import Chain
from langchain_community.chat_models import ChatOpenAI
from langchain_core.callbacks import CallbackManagerForChainRun
from langchain_core.language_models import BaseLanguageModel
from langchain_core.prompts import BasePromptTemplate, PromptTemplate

import __init__


class wiki_article_chain(Chain):
    """开发一个wiki文章生成器"""
    prompt: BasePromptTemplate
    llm: BaseLanguageModel
    out_key: str = "text"

    @property
    def input_keys(self) -> List[str]:
        """将返回Prompt所需的所有键"""
        return self.prompt.input_variables

    @property
    def output_keys(self) -> List[str]:
        """将始终返回text键"""
        return [self.out_key]

    def _call(
            self,
            inputs: Dict[str, Any],
            run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        """运行链"""
        prompt_value = self.prompt.format_prompt(**inputs)
        print("prompt_value:", prompt_value)
        response = self.llm.generate_prompt(
            [prompt_value], callbacks=run_manager.get_child() if run_manager else None
        )
        print("response:", response)
        if run_manager:
            run_manager.on_text("wiki article is written")
        return {self.out_key: response.generations[0][0].text}

    @property
    def _chain_type(self) -> str:
        """链类型"""
        return "wiki_article_chain"


chain = wiki_article_chain(
    prompt=PromptTemplate(
        template="写一篇关于{topic}的维基百科形式的文章",
        input_variables=["topic"]
    ),
    llm=ChatOpenAI(
        temperature=0
    ),
)

result = chain.run({"topic": "降本增效"})
print(result)
